#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import PointCloud2
from nav_msgs.msg import OccupancyGrid
from visualization_msgs.msg import Marker, MarkerArray
import numpy as np
from scipy.ndimage import gaussian_filter
from geometry_msgs.msg import Pose
import tf2_ros
from tf2_ros import TransformException
from tf2_ros.buffer import Buffer
from tf2_ros.transform_listener import TransformListener
import sensor_msgs_py.point_cloud2 as pc2
import math
from rclpy.qos import QoSProfile, ReliabilityPolicy

class PointCloudToCostmap(Node):
    def __init__(self):
        super().__init__('pointcloud_to_costmap')

        qos_profile = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            depth=1
        )

        self.map_resolution = self.declare_parameter('resolution', 0.1).get_parameter_value().double_value
        self.inflation_radius = self.declare_parameter('inflation_radius', 0.3).get_parameter_value().double_value

        self.map_width = self.declare_parameter('width', 10.0).get_parameter_value().double_value
        self.map_height = self.declare_parameter('height', 10.0).get_parameter_value().double_value
        self.map_height_3d = self.declare_parameter('height_3d', 5.0).get_parameter_value().double_value
        self.pointcloud_z_threshold = self.declare_parameter('z_threshold', -0.3).get_parameter_value().double_value

        self.base_footprint_link = self.declare_parameter('base_footprint_link', 'base_footprint').get_parameter_value().string_value
        self.velodyne_link = self.declare_parameter('velodyne_link', 'velodyne').get_parameter_value().string_value

        self.points_raw_topic = self.declare_parameter('points_raw_topic', 'points_raw').get_parameter_value().string_value
        self.processed_pointcloud_topic = self.declare_parameter('processed_pointcloud_topic', 'processed_points').get_parameter_value().string_value
        
        self.obstacle_layer_topic = self.declare_parameter('obstacle_layer_topic', 'obstacle_layer').get_parameter_value().string_value
        self.inflation_layer_topic = self.declare_parameter('inflation_layer_topic', 'inflation_layer').get_parameter_value().string_value
        self.local_costmap_topic = self.declare_parameter('local_costmap_topic', 'local_costmap').get_parameter_value().string_value
        self.voxel_layer_topic = self.declare_parameter('voxel_layer_topic', 'voxel_layer').get_parameter_value().string_value
        
        self.pointcloud_sub = self.create_subscription(PointCloud2, self.points_raw_topic, self.pointcloud_callback, 10)
        self.obstacle_layer_pub = self.create_publisher(OccupancyGrid, self.obstacle_layer_topic, 10)
        self.inflation_layer_pub = self.create_publisher(OccupancyGrid, self.inflation_layer_topic, 10)
        self.local_costmap_pub = self.create_publisher(OccupancyGrid, self.local_costmap_topic, 10)
        self.voxel_marker_pub = self.create_publisher(MarkerArray, self.voxel_layer_topic + '_markers', 10)
        self.voxel_grid_pub = self.create_publisher(OccupancyGrid, self.voxel_layer_topic, 10)
        self.processed_pointcloud_pub = self.create_publisher(PointCloud2, self.processed_pointcloud_topic, 10)
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)

        width_cells = int(self.map_width / self.map_resolution)
        height_cells = int(self.map_height / self.map_resolution)
        height_cells_3d = int(self.map_height_3d / self.map_resolution)
        self.obstacle_layer = np.zeros((height_cells, width_cells), dtype=np.int8)
        self.inflation_layer = np.zeros_like(self.obstacle_layer)
        self.voxel_layer = np.zeros((height_cells_3d, height_cells, width_cells), dtype=np.int8)

    def pointcloud_callback(self, msg):
        try:
            trans = self.tf_buffer.lookup_transform(self.velodyne_link, msg.header.frame_id, rclpy.time.Time())
            translation = (trans.transform.translation.x, trans.transform.translation.y, trans.transform.translation.z)
            rotation = (trans.transform.rotation.x, trans.transform.rotation.y, trans.transform.rotation.z, trans.transform.rotation.w)
        except TransformException as ex:
            self.get_logger().warn(f"Could not transform {msg.header.frame_id} to {self.velodyne_link}: {ex}")
            return

        points = self.transform_pointcloud(msg, translation, rotation)
        filtered_points = self.remove_ground_points(points, self.pointcloud_z_threshold)
        processed_cloud = self.create_pointcloud2(msg.header, filtered_points)
        self.processed_pointcloud_pub.publish(processed_cloud)

        obstacle_map = self.generate_obstacle_layer(filtered_points, translation)
        inflated_map = self.inflate_obstacle_layer(obstacle_map)
        voxel_map = self.generate_voxel_layer(filtered_points, translation)

        self.publish_costmap(obstacle_map, translation, self.obstacle_layer_pub)
        self.publish_costmap(inflated_map, translation, self.inflation_layer_pub)
        self.publish_voxel_markers(voxel_map, translation)
        self.publish_voxel_grid(voxel_map, translation)

        local_costmap = np.clip(obstacle_map + inflated_map, 0, 100).astype(np.int8)
        self.publish_local_costmap(local_costmap, translation)

    def transform_pointcloud(self, cloud, translation, rotation):
        trans_matrix = np.eye(4)
        trans_matrix[:3, 3] = np.array(translation)

        rot_matrix = self.quaternion_to_rotation_matrix(rotation)

        transform_matrix = np.dot(trans_matrix, rot_matrix)

        pc_data = list(pc2.read_points(cloud, field_names=("x", "y", "z"), skip_nans=True))

        if len(pc_data) == 0:
            raise ValueError("PointCloud2 数据为空或无法读取。")

        x_data = np.array([point[0] for point in pc_data], dtype=np.float32)
        y_data = np.array([point[1] for point in pc_data], dtype=np.float32)
        z_data = np.array([point[2] for point in pc_data], dtype=np.float32)
        points = np.vstack((x_data, y_data, z_data)).T

        if points.shape[1] != 3:
            raise ValueError(f"PointCloud2 数据形状不正确。期望 (N, 3)，但是得到: {points.shape}")

        transformed_points = np.dot(points, transform_matrix[:3, :3].T) + transform_matrix[:3, 3]

        return transformed_points

    def quaternion_to_rotation_matrix(self, quat):
        x, y, z, w = quat
        xx, yy, zz = x * x, y * y, z * z
        xy, xz, yz = x * y, x * z, y * z
        wx, wy, wz = w * x, w * y, w * z

        rot_matrix = np.array([
            [1 - 2 * (yy + zz), 2 * (xy - wz), 2 * (xz + wy), 0],
            [2 * (xy + wz), 1 - 2 * (xx + zz), 2 * (yz - wx), 0],
            [2 * (xz - wy), 2 * (yz + wx), 1 - 2 * (xx + yy), 0],
            [0, 0, 0, 1]
        ])
        return rot_matrix

    def remove_ground_points(self, points, z_threshold=-0.3):
        mask = points[:, 2] > z_threshold
        filtered_points = points[mask]
        return filtered_points

    def create_pointcloud2(self, header, points):
        cloud_msg = pc2.create_cloud_xyz32(header, points)
        return cloud_msg

    def generate_obstacle_layer(self, points, robot_translation, inflation_radius = 0.15):
        width_cells = int(self.map_width / self.map_resolution)
        height_cells = int(self.map_height / self.map_resolution)
        obstacle_layer = np.zeros((height_cells, width_cells), dtype=np.int8)

        for point in points:
            x_idx = int((point[0] - robot_translation[0] + self.map_width / 2) / self.map_resolution)
            y_idx = int((point[1] - robot_translation[1] + self.map_height / 2) / self.map_resolution)
            if 0 <= x_idx < width_cells and 0 <= y_idx < height_cells:
                obstacle_layer[y_idx, x_idx] = 100

        new_obstacle_layer = obstacle_layer.copy()

        for y in range(1, height_cells - 1):
            for x in range(1, width_cells - 1):
                if obstacle_layer[y, x] == 0:
                    if (obstacle_layer[y-1, x] == 100 and obstacle_layer[y+1, x] == 100 and 
                        obstacle_layer[y, x-1] == 100 and obstacle_layer[y, x+1] == 100):
                        new_obstacle_layer[y, x] = 100

        inflated_obstacle_layer = np.array(new_obstacle_layer)
        inflated_obstacle_layer = gaussian_filter(inflated_obstacle_layer, sigma=inflation_radius / self.map_resolution)
        inflated_obstacle_layer = np.clip(inflated_obstacle_layer, 0, 100)
        
        return inflated_obstacle_layer.astype(np.int8)

    def inflate_obstacle_layer(self, obstacle_layer):
        inflated_layer = gaussian_filter(obstacle_layer.astype(np.float32), sigma=self.inflation_radius / self.map_resolution)
        return np.clip(inflated_layer, 0, 100).astype(np.int8)

    def generate_voxel_layer(self, points, robot_translation):
        width_cells = int(self.map_width / self.map_resolution)
        height_cells = int(self.map_height / self.map_resolution)
        height_cells_3d = int(self.map_height_3d / self.map_resolution)
        
        voxel_layer = np.zeros((height_cells_3d, height_cells, width_cells), dtype=np.int8)

        for point in points:
            x_idx = int((point[0] - robot_translation[0] + self.map_width / 2) / self.map_resolution)
            y_idx = int((point[1] - robot_translation[1] + self.map_height / 2) / self.map_resolution)
            z_idx = int((point[2] + self.map_height_3d / 2) / self.map_resolution)

            if 0 <= x_idx < width_cells and 0 <= y_idx < height_cells and 0 <= z_idx < height_cells_3d:
                voxel_layer[z_idx, y_idx, x_idx] = 100

        return voxel_layer

    def publish_costmap(self, costmap, translation, pub):
        costmap_msg = OccupancyGrid()
        
        costmap_msg.header.frame_id = self.base_footprint_link
        costmap_msg.header.stamp = self.get_clock().now().to_msg()
        costmap_msg.info.resolution = self.map_resolution
        costmap_msg.info.width = 10
        costmap_msg.info.height = 10
        costmap_msg.info.origin.position.x = translation[0]
        costmap_msg.info.origin.position.y = translation[1]
        costmap_msg.info.origin.position.z = translation[2]

        costmap_clipped = np.clip(costmap, -128, 127).astype(np.int8)
        costmap_msg.data = costmap_clipped.ravel().tolist()

        pub.publish(costmap_msg)

    def publish_local_costmap(self, costmap, translation):
        costmap_msg = OccupancyGrid()
        costmap_msg.header.stamp = self.get_clock().now().to_msg()
        costmap_msg.header.frame_id = self.base_footprint_link
        costmap_msg.info.resolution = self.map_resolution
        costmap_msg.info.width = costmap.shape[1]
        costmap_msg.info.height = costmap.shape[0]
        costmap_msg.info.origin.position.x = translation[0] - self.map_width / 2
        costmap_msg.info.origin.position.y = translation[1] - self.map_height / 2

        costmap_clipped = np.clip(costmap, -128, 127).astype(np.int8)

        # 打印调试信息
        print("Costmap min:", costmap_clipped.min(), "max:", costmap_clipped.max())
        print("Costmap dtype:", costmap_clipped.dtype)

        costmap_msg.data = [int(value) for value in costmap_clipped.ravel()]

        if not all(-128 <= val <= 127 for val in costmap_msg.data):
            self.get_logger().error("Costmap data contains values out of the acceptable range [-128, 127].")
            return

        self.local_costmap_pub.publish(costmap_msg)


    def publish_voxel_markers(self, voxel_layer, translation):
        marker_array = MarkerArray()
        current_time = self.get_clock().now().to_msg()

        for z in range(voxel_layer.shape[0]):
            for y in range(voxel_layer.shape[1]):
                for x in range(voxel_layer.shape[2]):
                    if voxel_layer[z, y, x] > 0:
                        marker = Marker()
                        marker.header.stamp = current_time
                        marker.header.frame_id = self.base_footprint_link
                        marker.type = Marker.CUBE
                        marker.action = Marker.ADD
                        marker.pose.position.x = x * self.map_resolution + translation[0] - self.map_width / 2
                        marker.pose.position.y = y * self.map_resolution + translation[1] - self.map_height / 2
                        marker.pose.position.z = z * self.map_resolution - self.map_height_3d / 2
                        marker.pose.orientation.w = 1.0
                        marker.scale.x = self.map_resolution
                        marker.scale.y = self.map_resolution
                        marker.scale.z = self.map_resolution
                        marker.color.a = 0.8
                        marker.color.r = 1.0
                        marker.color.g = 0.0
                        marker.color.b = 0.0
                        marker_array.markers.append(marker)

        for i, marker in enumerate(marker_array.markers):
            marker.id = i

        self.voxel_marker_pub.publish(marker_array)

    def publish_voxel_grid(self, voxel_layer, translation):
        voxel_grid = np.max(voxel_layer, axis=0)
        self.publish_costmap(voxel_grid, translation, self.voxel_grid_pub)

def main(args=None):
    rclpy.init(args=args)
    node = PointCloudToCostmap()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
